"""
训练服务模块
"""

from datetime import datetime
from typing import List, Optional
from loguru import logger

from ..models.training_model import TrainingData, ModelVersion, TrainingJob


class TrainingService:
    """训练服务类"""

    def __init__(self):
        self.training_data: List[TrainingData] = []
        self.model_versions: List[ModelVersion] = []
        self.training_jobs: List[TrainingJob] = []

    def add_training_data(self, data: TrainingData) -> bool:
        """添加训练数据"""
        try:
            self.training_data.append(data)
            logger.info(f"训练数据添加成功: {data.id}")
            return True
        except Exception as e:
            logger.error(f"训练数据添加失败: {str(e)}")
            return False

    def get_training_data(self, data_id: int) -> Optional[TrainingData]:
        """获取训练数据"""
        for data in self.training_data:
            if data.id == data_id:
                return data
        return None

    def create_training_job(self, job: TrainingJob) -> bool:
        """创建训练任务"""
        try:
            self.training_jobs.append(job)
            logger.info(f"训练任务创建成功: {job.id}")
            return True
        except Exception as e:
            logger.error(f"训练任务创建失败: {str(e)}")
            return False

    def get_training_job(self, job_id: int) -> Optional[TrainingJob]:
        """获取训练任务"""
        for job in self.training_jobs:
            if job.id == job_id:
                return job
        return None 